"""
Convience functions for working with various frequently used itk filters
"""
import itk
import numpy as na
import scipy.ndimage as ndi
import numpy.numarray.mlab as mlab

import sys
import os
import os.path
from . import query
#import waterMark

def getMattesMetric( iType1, iType2, numBins = 80, 
                    numSamples = 100000, verbose = True ):
    """getMattesMetric: generates a Mattes Metric for mutual information based
    registration.
    iType1: itk image type of target image
    iType2: itk image typye of source image
    numBins: number of bins to use in the histogram
    numSamples: number of image samples from which to generate the histogram
    verbose: provide deailed output if True"""
    try:
        if( verbose):
            print iType1,iType2
        imageMetric  = itk.MattesMutualInformationImageToImageMetric[iType1,iType2].New()
        if( verbose):
            print "Setting number of bins to use"
        imageMetric.SetNumberOfHistogramBins( numBins )
        if( verbose ):
            print "Setting number of spatial samples to use"
            print "numSamples=%d",numSamples
        imageMetric.SetNumberOfSpatialSamples( numSamples )
        
        return imageMetric
    except Exception, error:
        print "failed in getMattestMetric", error
def getMeanSquareMetric(iType1,iType2 ):
    """getMeanSquareMetric: create a mean square image to image metric for the
    specified target and source image types"""
    try:
        imageMetric  = itk.MeanSquaresImageToImageMetric[iType1,iType2].New()
        return imageMetric
    except Exception, error:
        print "failed in getMeanSquareMetric", error



def mergeArraysToRGB(r,g,b,intype=itk.Image[itk.F,2]):
    """takes three itk images and merges them into a single PIL RGB image"""
    try:
        from ..imageUtils.basicImageUtils import  ArrayToPilImageStraight
        ra = itk.PyBuffer[intype].GetArrayFromImage(r)
        ga = itk.PyBuffer[intype].GetArrayFromImage(g)
        ba = itk.PyBuffer[intype].GetArrayFromImage(b)
        ri = ArrayToPilImageStraight(ra.astype(na.uint8))
        gi = ArrayToPilImageStraight(ga.astype(na.uint8))
        bi = ArrayToPilImageStraight(ba.astype(na.uint8))
        img = Image.merge('RGB',(ri,gi,bi))
        return img
    except Exception, error:
        print "failed in mergeArraysToRGB", error
        return None

def getMaskExtent(mask,iType=itk.Image[itk.UC,3]):
    """Given an ITK image use python to determine the extent of the nonzero regions"""
    try:
        from ..imageUtils.basicImageUtils import get_crds
        array = itk.PyBuffer[iType].GetArrayFromImage(mask)
        inds = na.nonzero(array.flat)[0]
        crds = get_crds(inds,array.shape)
        xrange = ((crds[0]).min(),(crds[0]).max())
        yrange = ((crds[1]).min(),(crds[1]).max())
        zrange = ((crds[2]).min(),(crds[2]).max())
        print "mask extent is", xrange, yrange, zrange
        return xrange,yrange,zrange, array.shape[::-1] # return shape with ITK ordering
    except Exception, error:
        print "failed in getMaskExtent", error
            
def subRegionZFromMask(img, mask, iType = itk.Image[itk.F,3], mType=itk.Image[itk.UC,3]):
    """Automatically subregions an image to contain only the z region including non-zero
    mask values"""
    try:
        print "getting mask extent"
        bb = getMaskExtent(mask,mType)
        
        minz = max(bb[2][0]-5,0)
        maxz = min(bb[2][1]+5,bb[-1][2]-1)
        
        index = [0,0,minz]
        sz = [bb[-1][0],bb[-1][1],maxz-minz+1]
        print "subregionging with index %s and size %s"%(index,sz)
        newImg = autoSubRegion(img,index,sz,iType)
        newMask = autoSubRegion(mask,index,sz,mType)
        return newImg, newMask
    except Exception, error:
        print "failed in subRegionZFromMask", error
        
   
def autoSubRegion(img,index,sz, iType=itk.Image[itk.F,3]):
    """extract a subregion from img defined by origin and sz"""
    try:
        subFilter = itk.RegionOfInterestImageFilter[iType,iType].New()
        subFilter.SetInput(img)
        dr = subFilter.GetRegionOfInterest()
        dr.SetIndex(index)
        dr.SetSize(sz)
        subFilter.SetRegionOfInterest(dr)
        img2 = subFilter.GetOutput()
        img2.Update()
        
        # Get origin, spacing and direction. April 11, 2007: I'm assuming (dangerously)
        # for the moment that image is increasing in z from origin
        origin = img.GetOrigin()
        spacing = img.GetSpacing()
        for i in range(origin.GetPointDimension()):
            newOrigin = origin.GetElement(i)+index[i]*spacing.GetElement(i)
            origin.SetElement(i,newOrigin)
            
        img2.SetOrigin(origin)
        img2.SetSpacing(img.GetSpacing())
        img2.SetDirection(img.GetDirection())
        return img2
    except Exception, error:
        print "failed in autoSubRegion()", error
        return None
        
                    
def subRegionImage(img, iOrigin=None, iSize=None,
                    modOrigin=None):
    """subregion itk image based on specified origin and size
    img: an itk image
    iOrigin: (minx,miny,minz) location of subregion
    iSize: (deltax, delty, deltz)
    modOrigin: ???"""
    try:
        print "Incoming range",itk.range(img)
        inIType = query.getImageType(img)
        subFilter = itk.RegionOfInterestImageFilter[inIType,inIType].New()
        subFilter.SetInput(img)
        ir = subFilter.GetRegionOfInterest()
        index = ir.GetIndex()
        size = ir.GetSize()
        print "Setting origin"
        if( iOrigin == None ):
            iOrigin = input("Enter tuple of %s origin for subregion\n"%prompt)
        for i in range(len(iOrigin)):
            index.SetElement(i,iOrigin[i])
        print "Setting size"
        if( iSize == None ):
            iSize = input("Enter tuple of %s size for subregion\n"%prompt)
        for i in range(len(iSize)):
            size.SetElement(i,iSize[i])

        ir.SetIndex(index)
        ir.SetSize(size)
        subFilter.SetRegionOfInterest(ir)
        img2 = subFilter.GetOutput()
        print "Outgoing range",itk.range(img2)
        if( modOrigin != None ):
            print "modifying origin"
            try:
                img2.SetOrigin(modOrigin)
            except:
                try:
                    io = img2.GetOrigin()
                    for i in range(len(modOrigin)):
                        io.SetElement(i,modOrigin[i])
                except:
                    print "Could not reset origin"

        return img2
    except Exception, error:
        print "failed in subRegionImage", error
        return img 

def maskImage(img, mask=None, itype=itk.F,mtype=itk.UC,
              fname=None, maskValue = 1, minv=None, closeValue = 0, 
              record=None):
    """maskImage: multiply an image by a binary mask. The mask is either passed
    in as an argument or a filename pointing to the mask is provided
    img: image to mask
    mask: argument to pass mask to function with
    fname: name of file containing mask
    
    NOTE: mask overrides fname"""
    try:
        if( fname == None and mask == None ):
            print "did not provide a filename for mask or an existing mask"
            return img, None
        dim = query.getImgDim(img)    

        iitype = itk.Image[itype,dim]
        mitype = itk.Image[mtype,dim]
        if( mask == None ):
            print "Reading mask from ", fname
            print "creating reader", mtype, dim
            maskReader = itk.ImageFileReader[mitype].New(FileName=fname)
            mask = maskReader.GetOutput()      

        print "getting arrays"
        marray = na.array(itk.PyBuffer[mitype].GetArrayFromImage(mask))
        farray = na.array(itk.PyBuffer[iitype].GetArrayFromImage(img))


        # Get the min value present in farray. For CT this will
        # be the regions where an image was not reconstructed
        if( minv == None ):
            minv,maxv = itk.range(img)

        # create a mask of the maskValue locations
        print "create maskvalue mask"
        m = na.where(marray == maskValue,1,0)
        if( closeValue ):
            print "marray type is ",m.typecode(),m.shape
            
            m = ndi.binary_closing(m,iterations=closeValue)
            print "marray type is ",m.typecode(),m.shape
            temp = 0
        print "get nonzero inds"
        inds = na.nonzero(m.flat != 1)[0]
        print "put the values"
        print m.shape, farray.shape, inds.shape, minv
        farray.put(inds,minv)
        img2 = itk.PyBuffer[iitype].GetImageFromArray(farray)
        
        img2.SetSpacing(img.GetSpacing())
        img2.SetDirection(img.GetDirection())
        img2.SetOrigin(img.GetOrigin())
        
        if( not record ):
            record = './temp.mha'
        writer = itk.ImageFileWriter[iitype].New(FileName=record)
        writer.SetInput(img2)
        writer.Update()
        reader = itk.ImageFileReader[iitype].New(FileName=record)
        img2 = reader.GetOutput()
        img2.Update()
        os.remove(record)
        
        return img2
    except Exception, error:
        print "failed in maskImage", error
        return img, farray, inds, minv
def maskImage2(img,itype=itk.F,mtype=itk.UC,fname=None, 
               maskValue = 1, minv=None, closeValue = 0, record=None):
    """an older version of maskImage that requires the mask be pased via
    filename"""
    try:
        if( fname == None ):
            print "did not provide a filename for mask"
            return img
        
        dim = query.getImgDim(img)   
        iitype = itk.Image[itype,dim]
        mitype = itk.Image[mtype,dim]
        print "creating reader", mtype, dim
        maskReader = itk.ImageFileReader[mitype].New(FileName=fname)
        maskImage = maskReader.GetOutput()
        maskImage.Update()
        printImageData(maskImage,"maskImage")
        mask = itk.ImageMaskSpatialObject._3.New()
        mask.SetImage(maskImage)
        mask.Update()
        return mask

    except Exception, error:
        print "failed in maskImage2", error
        return None

def getFlipFilter(img,axis=0,mitype=None, mtype=itk.UC, dim=3 ):
    """This is currently misnamed. This flips the image; it does not return the
    flip filter"""
    try:
        if( not mitype ):
            mitype = itk.Image[mtype,dim]   
        print mitype
        flip = itk.FlipImageFilter[mitype].New()
        print "flipping Image"
        flip.SetInput(img)
        t0 = flip.GetFlipAxes()
        t0.SetElement(axis,True)
        flip.SetFlipAxes(t0)
        flip.Update()
        return flip

    except Exception, error:
        print "failed in flipImage", error
        return None        

def flipImage(img,axis=0,mitype=None, mtype=None, imgMode='sshort',
              dim=3, **args):
    """flipImage: flip an image using ITK
    axis: x,y,z axis to flip image around
    mitype:???
    mtype: ???
    """
    try:
        if( not mitype ):
            if( not mtype ):
                mtype = query.getITKType(imgMode)
            mitype = itk.Image[mtype,dim]   
        print mitype
        flip = itk.FlipImageFilter[mitype].New()
        print "flipping Image"
        flip.SetInput(img)
        t0 = flip.GetFlipAxes()
        t0.SetElement(axis,True)
        flip.SetFlipAxes(t0)
        flip.Update()
        img2 = flip.GetOutput()
        img2.Update()

        return img2

    except Exception, error:
        print "failed in flipImage", error
        return img        
def maskImage3(img,itype=itk.F,mtype=itk.UC,fname=None, mode = 'numarray',flipMask=-1, **args):
    """Subregions img and mask based on bounding region of specified mask. Returns 
    ImageMaskSpatialObject corresponding to masked image for use in registration optimization"""
    try:
        print "maskImage3"
        if( fname == None ):
            print "did not provide a filename for mask"
            return img
        
        print "Reading mask from ", fname
        dim = query.getImgDim(img)   
        iitype = itk.Image[itype,dim]
        mitype = itk.Image[mtype,dim]
        print "creating reader", mtype, dim
        maskReader = itk.ImageFileReader[mitype].New(FileName=fname)
        maskImg = maskReader.GetOutput()
        maskImg.Update()
        if( flipMask >= 0 ):
            maskImg = flipImage(maskImg,axis=flipMask,mtype=mtype, dim=dim)
        print "Subregioning images"
        img2,maskImg =  subRegionZFromMask(img, maskImg, iType = iitype, mType=mitype)

        query.printImageData(maskImg,"maskImage")
        if( mode == 'numarray' ):
            print "multiplying image by mask"
            img2 = maskImage(img2,mask=maskImg,itype=itype,mtype=mtype)
            return img2, maskImg
        else:
            mask = itk.ImageMaskSpatialObject._3.New()
            mask.SetImage(maskImage)
            mask.Update()
            return img2, mask

    except Exception, error:
        print "failed in maskImage3", error
        return None

def castImage(source,oitype=itk.Image[itk.F,3]):
    """typecast an ITK image"""
    inType = query.getImgPixelDimType(source)
    print "castImage",oitype
    caster = getCaster(source,iitype = inType, outType = oitype)
    oimg = caster.GetOutput()
    oimg.Update()
    return oimg

def getCaster(source, outType = None, **kwargs):
    """Create a caster for an ITK image"""
    print "source type is",type(source),"out type is",outType
    inType = query.getImageType(source)
    try:
        minv,maxv = itk.range(source.GetOutput())
    except:
        print "behaving like image"
        minv,maxv = itk.range(source)
    print minv, maxv
    caster =itk.RescaleIntensityImageFilter[inType,outType].New()
    caster.SetOutputMinimum( int(minv) )
    caster.SetOutputMaximum( int(maxv) )
    try:
        caster.SetInput(source.GetOutput())
    except:
        caster.SetInput(source)
    return caster

def rescaleImage(img,scale=None):
    """rescale an ITK image to that specified by scale"""
    if( scale == None ):
        scale = (0,1000)
    iitype = query.getImgPixelDimType(img)
    caster = itk.RescaleIntensityImageFilter[iitype,iitype].New()
    caster.SetOutputMinimum(scale[0])
    caster.SetOutputMaximum(scale[1])
    caster.SetInput(img)
    oimg = caster.GetOutput()
    oimg.Update()
    return oimg


def getImageCaster(cs,itype=itk.F,otype=itk.SS):
    try: 
        dim = query.getObjectSize(cs)
        iitype = itk.Image[itype,dim]
        oitype = itk.Image[otype,dim]
        minv,maxv = itk.range(cs)
        print minv, maxv
        caster =itk.RescaleIntensityImageFilter[iitype,oitype].New()
        caster.SetOutputMinimum( int(minv) )
        caster.SetOutputMaximum( int(maxv) )
        caster.SetInput(cs)
        return caster
    except Exception, error:
        print "failed in getImageCaster", error
        return None 
def smoothImage(img,sigma, imgType = None, inType = itk.F, numDim = 3 ):
    """apply Gaussian smoothing to image independently in each direction"""
    try:
        if( imgType == None ):
                imgType = itk.Image[inType,numDim]
                
        if( type(sigma) == type(1) or type(sigma) == type(1.0) ):
            sigma = [sigma]*numDim
        smootherX = itk.RecursiveGaussianImageFilter[imgType,imgType].New()
        smootherY = itk.RecursiveGaussianImageFilter[imgType,imgType].New()
        smootherX.SetInput(img)
        smootherY.SetInput(smootherX.GetOutput())
        smootherX.SetSigma(sigma[0])
        smootherY.SetSigma(sigma[1])
        smootherX.SetDirection(0)
        smootherY.SetDirection(1)
        smootherX.SetNormalizeAcrossScale(False)
        smootherY.SetNormalizeAcrossScale(False)
        if( numDim == 3 ):
            smootherZ = itk.RecursiveGaussianImageFilter[imgType,imgType].New()
            smootherZ.SetInput(smootherY.GetOutput())
            smootherZ.SetSigma(sigma[2])
            smootherZ.SetDirection(2)
            smootherZ.SetNormalizeAcrossScale(False)
            img2 =  smootherZ.GetOutput()
            img2.Update()
        else:
            img2 = smootherY.GetOutput()
            img2.Update()
        return img2
    except Exception, error:
        print "failed in smoothImge", error        
                            

        
def subsampleImage(img, samp, 
                    debug = False, **kwargs):
    try:
        inputSize = img.GetLargestPossibleRegion().GetSize()
        imgIType = query.getImageType(img)
        
        if( debug ): print "creating resampler"
        resampler = itk.ResampleImageFilter[imgIType,imgIType].New()
        dim = imgIType.GetImageDimension()

        inputSize = img.GetLargestPossibleRegion().GetSize()

        resampler.SetInput( img )
        print "set input"
        inputSpacing = img.GetSpacing()
        inputSize = img.GetLargestPossibleRegion().GetSize()
        interpolator = itk.LinearInterpolateImageFunction[imgIType,itk.D].New()
        transform = itk.IdentityTransform[itk.D,imgIType.GetImageDimension()].New()
        
        outputSize = []
        outputSpacing = []
        for i in range( dim ):   
                outputSize.append(int( inputSize.GetElement(i)/float(samp[i]) ) )
                outputSpacing.append(inputSpacing.GetElement(i)*samp[i])
        
        
        if( debug ): print "setting interpolator"
        transform.SetIdentity()
        
        resampler.SetInterpolator(interpolator)
        resampler.SetTransform(transform)
        
        
        if( debug ): print "setting output spacing and size"
        resampler.SetOutputSpacing(outputSpacing)
        resampler.SetOutputOrigin(img.GetOrigin())
        resampler.SetSize(outputSize)
        img2 = resampler.GetOutput()

        img2.Update()
        return img2

    except Exception, error:
        print "failed in subsampleImage", error
def subsample2DImage(img, samp, inType = itk.F, smooth = False):
    try:
        resampler = itk.ResampleImageFilter[itk.Image[itk.F,2],itk.Image[itk.F,2]].New()
        dim = query.getObjectSize(img)
        inputSize = img.GetLargestPossibleRegion().GetSize()
        caster = itk.CastImageFilter[itk.Image[inType,2],itk.Image[itk.F,2]].New()
        caster.SetInput(img)
        img2 = caster.GetOutput()
        img2.Update()
        if( smooth ):
            smootherX = itk.RecursiveGaussianImageFilter[itk.Image[itk.F,2],itk.Image[itk.F,2]].New()
            smootherY = itk.RecursiveGaussianImageFilter[itk.Image[itk.F,2],itk.Image[itk.F,2]].New()
            smootherX.SetSigma(samp) # Add image spacing correction
            smootherY.SetSigma(samp) # Add image spacing correction
            smootherX.SetInput(img2)
            smootherY.SetInput(smootherX.GetOutput())
            smootherX.SetDirection(0)
            smootherY.SetDirection(1)
            smootherX.SetNormalizeAcrossScale(False)
            smootherY.SetNormalizeAcrossScale(False)
            resampler.SetInput( smootherY.GetOutput() )
        else:
            resampler.SetInput( img2 )
        transform = itk.IdentityTransform.D2.New()
        transform.SetIdentity()
        interpolator = itk.LinearInterpolateImageFunction.IF2D.New()
        resampler.SetInterpolator(interpolator)
        resampler.SetTransform(transform)
        
        inputSpacing = img.GetSpacing()
        outputSpacing = [inputSpacing.GetElement(0)*samp,inputSpacing.GetElement(1)*samp]
        #outputSpacing = [inputSpacing.GetElement(0),inputSpacing.GetElement(1)]
        print "outputSpacing=",outputSpacing
        resampler.SetOutputSpacing(outputSpacing)
        resampler.SetOutputOrigin(img.GetOrigin())
        
        inputSize = img.GetLargestPossibleRegion().GetSize()
        outputSize = [int( inputSize.GetElement(0)/float(samp) ), int(inputSize.GetElement(1)/float(samp) )]
        print 'setting outputsize'
        resampler.SetSize(outputSize)
        caster2 = itk.CastImageFilter[itk.Image[itk.F,2],itk.Image[inType,2]].New()
        caster2.SetInput(resampler.GetOutput())
        print 'getting output'
        img3 = caster2.GetOutput()
        img3.Update()
        return img3
    except Exception, error:
        print "failed in subsampleImage", error
                
                
def thresholdImage(img, range, map, **kwargs ): 
    """Creates a binary image from img pixels lying within range[0] and range[1] are assigned
    a value map[1]. Pixels/voxels lying outside of this range are assinged a value map[0]"""
    try:

        iType = query.getImageType(img)
        dim = iType.GetImageDimension()
        ft = itk.BinaryThresholdImageFilter[iType,itk.Image[itk.UC,dim]].New()
        
        ft.SetInput(img)
        ft.SetOutsideValue(map[0])
        ft.SetInsideValue(map[1])
        ft.SetLowerThreshold(range[0])
        ft.SetUpperThreshold(range[1])
        img2 = ft.GetOutput()
        img2.Update()
        return img2
    except Exception, error:
        print "failed in thresholdImage", error
        return None
        

def getImageCenterOfGravity(img, iType, getAxes = False):
    """returns the image center-of-gravity and principal axes (if requested) of passed image img"""
    try:
        mc = itk.ImageMomentsCalculator[iType].New()
        mc.SetImage(img)
        mc.Compute()
        cg = mc.GetCenterOfGravity()
        print "center of gravity [%f,%f]"%(cg.GetElement(0),cg.GetElement(1))
        if( getAxes ):
            pa = mc.GetPrincipalAxes()
            pa2 = pa.GetVnlMatrix()
            return cg, pa2
        else:
            return cg
    except Exception, error:
        print "failed in getImageCenterOfGravity()", error

def getTransform2D(options, fixedImage, movingImage, registration, iType = itk.Image[itk.F,2], useCenterOfGravity = True):
    """Defines and initializes a 2D registration between fixedImage and movingImage"""
    try:
        
        if( useCenterOfGravity ): # If we are going to use moments to initialize transform compute them
            print "using center of gravity to initialize transformation"
            print_itk_variable(movingImage.GetOrigin(),name="Moving Image Origin")
            print_itk_variable(fixedImage.GetOrigin(), name="Fixed Image Origin")  
            print_itkVector(movingImage.GetSpacing(), name="moving image spacing")
            print_itkVector(fixedImage.GetSpacing(), name="fixed image spacing")
            fcg, fpa = getImageCenterOfGravity(fixedImage, iType, getAxes = True)
            mcg, mpa = getImageCenterOfGravity(movingImage, iType, getAxes = True )
            ix = -1.0*float(fcg.GetElement(0)-mcg.GetElement(0))
            iy = -1.0*float(fcg.GetElement(1)-mcg.GetElement(1))
            #raw_input('continue')
        
        if( options.transform == 'translation' ):
            transform    = itk.TranslationTransform.D2.New()
            initialParameters = transform.GetParameters()
            initialParameters.SetElement(0,0.)
            if( not useCenterOfGravity ):
                ix = 0.0
                iy = 0.0
            if( options.initTranslation ):
                ix = float(options.initTranslation[0])
                iy = float(options.initTranslation[1])
            print "initial transformation is estimated as [%f,%f]"%(ix,iy)
            registration.SetInitialTransformParameters( [ix,iy] )
            
            scales = [1/1000.0]*2
        else:
            print "Using centered transform"
            transform = itk.CenteredRigid2DTransform.D.New()
            # Initial transform parameters 
            #
#             print "Setting initial angle to %f"%(0.0)
#             if( useCenterOfGravity ):
#                 f1 = fpa.get_row(1)
#                 m1 = mpa.get_row(1)
#                 f1v = na.array((f1.get(0),f1.get(1)))
#                 m1v = na.array((m1.get(0),m1.get(1)))
#                 print f1v,m1v
#                 theta = na.arccos( na.dot(f1v,m1v)/na.sqrt(na.dot(f1v,f1v)*na.dot(m1v,m1v)))
#             else:
#                 theta = 0.0
            theta = 0.0
            print "setting angle to %f"%theta
            #raw_input('continue')
            transform.SetAngle( theta );
            

            # transform center
            if( useCenterOfGravity ):
                print "Setting transform center of rotation"
                transform.SetCenter([fcg.GetElement(0),fcg.GetElement(1)])
                print "setting initial translation"
                transform.SetTranslation([ix,iy])
            else:
                # center of the fixed image
                fixedSpacing = fixedImage.GetSpacing()
                fixedOrigin = fixedImage.GetOrigin()
                fixedSize = fixedImage.GetLargestPossibleRegion().GetSize()
                
                centerFixed = ( fixedOrigin.GetElement(0) + fixedSpacing.GetElement(0) * fixedSize.GetElement(0) / 2.0,
                                fixedOrigin.GetElement(1) + fixedSpacing.GetElement(1) * fixedSize.GetElement(1) / 2.0 )
                
                # center of the moving image 
                movingSpacing = movingImage.GetSpacing()
                movingOrigin = movingImage.GetOrigin()
                movingSize = movingImage.GetLargestPossibleRegion().GetSize()
                
                centerMoving = ( movingOrigin.GetElement(0) + movingSpacing.GetElement(0) * movingSize.GetElement(0) / 2.0,
                                    movingOrigin.GetElement(1) + movingSpacing.GetElement(1) * movingSize.GetElement(1) / 2.0  )
                center = transform.GetCenter()
                center.SetElement( 0, centerFixed[0] )
                center.SetElement( 1, centerFixed[1] )
                transform.SetCenter(center)
                
                # transform translation
                translation = transform.GetTranslation()
                translation.SetElement( 0, centerMoving[0] - centerFixed[0] )
                translation.SetElement( 1, centerMoving[1] - centerFixed[1] )
                transform.SetTranslation(translation)
                
            print "setting registration to initial transform parameters"
            initialParameters = transform.GetParameters()
            
            print "Initial Parameters: "
            print "Angle: %f" % (initialParameters.GetElement(0), )
            print "Center: %f, %f" % ( initialParameters.GetElement(1), initialParameters.GetElement(2) )
            print "Translation: %f, %f" % (initialParameters.GetElement(3), initialParameters.GetElement(4))
            
            registration.SetInitialTransformParameters( initialParameters )  
            
            # Now we need to set up the transform Scales for the optimizer
            
            transform.GetNumberOfParameters()
            translationScale = 1.0/1000.0
            scales = [5]+[100]*2+[translationScale]*2
        print "initial scales", scales
        return transform, scales
    except Exception, error:
        print "failed in getTransform2D", error
        sys.exit(0)

def getScaleLogarithmicTransform(options, fixedImage, movingImage, registration, iType = itk.Image[itk.F,3]):
    """Defines and initializes a ScaleLogarithmicTransform between fixedImage and movingImage"""
    try:
        # Define transform
        transform = itk.ScaleLogarithmicTransform.D3.New()
        registration.SetInitialTransformParameters(transform.GetParameters())
        translationScale = 1.0/1000.0
        scales = [10000.0,10000.0,translationScale]
        stepSizes = [0.0001,0.01]
        print "initial scales", scales
        return transform, scales, stepSizes
    except Exception, error:
        print "failed in getScaleLogarithmicTransform", error
        sys.exit(0) 
def getVersorTransform3D(options, fixedImage, movingImage, registration, iType = itk.Image[itk.F,3]):
    """Defines and initializes a 3D versor registration between fixedImage and movingImage"""
    try:
        # Define transform
        transform = itk.VersorRigid3DTransform.D.New()
        
        # Now define initializer
        print iType
        initializer = itk.CenteredVersorTransformInitializer[iType,iType].New()
        initializer.SetMovingImage(movingImage)
        initializer.SetFixedImage(fixedImage)
        initializer.SetTransform(transform)
        
        # initialize center and translation
        initializer.MomentsOn()
        initializer.InitializeTransform()
        v1 = transform.GetVersor()
        v1.SetRotationAroundX(0.0)
        v1.SetRotationAroundY(0.0)
        v1.SetRotationAroundZ(0.0)
        transform.SetRotation(v1)
        
        registration.SetInitialTransformParameters(transform.GetParameters())
        transform.GetNumberOfParameters()
        stepSizes =  [0.0001,0.2]
        translationScale = 1.0/1000.0
        scales = [1]*3+[translationScale]*3
        print "initial scales", scales
        return transform, scales, stepSizes
    except Exception, error:
        print "failed in getVersorTransform3D", error
        sys.exit(0)          
def getMetric( iType,numPixels,metric='mattes', useAllSpatialSamples=False,numBins=100, spatialSamples=0.3 ):
    """based on values passed, return a requested image registation metric"""
    try:
        print metric, numBins, spatialSamples, numPixels
        if( metric == 'mattes' ):
            print "using mattes metric"
            imageMetric = getMattesMetric(iType,iType,numBins,int(spatialSamples*numPixels))
            if( useAllSpatialSamples ):
                imageMetric.UseAllPixelsOn()
                print imageMetric.GetUseAllPixels()
        elif( metric == 'kappa' ):
            print "using kappa metric"
            imageMetric =  itk.KappaStatisticImageToImageMetric[iType,iType].New()      
        else:
            print "using mean square metric"
            imageMetric = getMeanSquareMetric(iType,iType)
        #raw_input("continue")
        return imageMetric
    except Exception, error:
        print "failed in getMetric()", error

def getMetricOld( options, iType, numPixels ):
    """based on values stored in options, return a requested image registation metric"""
    try:
        print options.metric, options.numBins, options.spatialSamples, numPixels
        if( options.metric == 'mattes' ):
            print "using mattes metric"
            imageMetric = getMattesMetric(iType,iType,options.numBins,int(options.spatialSamples*numPixels))
            if( options.useAllSpatialSamples ):
                imageMetric.UseAllPixelsOn()
                print imageMetric.GetUseAllPixels()
        elif( options.metric == 'kappa' ):
            print "using kappa metric"
            imageMetric =  itk.KappaStatisticImageToImageMetric[iType,iType].New()      
        else:
            print "using mean square metric"
            imageMetric = getMeanSquareMetric(iType,iType)
        #raw_input("continue")
        return imageMetric
    except Exception, error:
        print "failed in getMetric()", error
        raw_input("continue")
        sys.exit(0)
def getMaximumImage(img1,img2):
    """Return an image that is the pixelwise maximum of img1 and img2"""
    try:
        iType1 = query.getImageType(img1)
        iType2 = query.getImageType(img2)
        maxFilter = itk.MaximumImageFilter[iType1,iType2,iType1].New()
        maxFilter.SetInput1(img1)
        maxFilter.SetInput2(img2)
        maxImg = maxFilter.GetOutput()
        maxImg.Update()
        return maxImg
    except Exception, error:
        print "could not return maximum image", error
def getRegistration( regScale, numLevels, iType, fixedImage2,  movingImage2):
    try:
        if( regScale == 'single' ):  
            print "Using single-scale registration" 
            registration = itk.ImageRegistrationMethod[iType,iType].New()

            
        else: # Assume multiscale
            # Get Image Pyramids
            print "Using multiscale registration with %d scales"%numLevels
            fixedImagePyramid = itk.RecursiveMultiResolutionPyramidImageFilter[iType,iType].New()
            movingImagePyramid = itk.RecursiveMultiResolutionPyramidImageFilter[iType,iType].New()
            fixedImagePyramid.SetInput(fixedImage2)
            movingImagePyramid.SetInput(movingImage2)
            registration = itk.MultiResolutionImageRegistrationMethod[iType,iType].New()
            registration.SetFixedImagePyramid( fixedImagePyramid )
            registration.SetMovingImagePyramid( movingImagePyramid ) 
            registration.SetNumberOfLevels(numLevels)
            
        registration.SetFixedImage(  fixedImage2  )
        registration.SetMovingImage( movingImage2 )
        registration.SetFixedImageRegion( fixedImage2.GetBufferedRegion() )

        return registration
    except Exception, error:
        print "failed in getRegistration", error
        
def transformImage(resampler,img=None,iType=itk.Image[itk.F,3]):
    """given the transform defined in registration, transform img in img or located on disk in fname
    and save to disk"""
    try:
        if( img ):
            resampler.SetInput(img)
            
        oimg = resampler.GetOutput()
        oimg.Update()
        return oimg
    except Exception, error:
        print "failed in transformImage", error
def interpolateImage(img, samp=None, interp = 'linear', debug=False,
                     returnSampling = False):
    """Takes an image and interplates based on factors assinged in samp. If samp is none, it is assumed that we are interplating to an isotropic voxel size"""
    try:
        # Get image type and dimension
        if( debug ): print "determining image type"
        imgIType = query.getImageType(img)
        numDim = imgIType.GetImageDimension()
        if( debug ): print "inType=%s"%(imgIType)

        inputSize = img.GetLargestPossibleRegion().GetSize()
        # set up constants to use throughout function
        
        
        if( debug ): print "creating resampler"
        resampler = itk.ResampleImageFilter[imgIType,imgIType].New()
        resampler.SetInput(img)
        if( samp == None ): # If no sampling is specified, assume interpolate to isotropic
            if( debug ): print "determining spacing and interpolation factors"
            spacing = query.get_itkVector(img.GetSpacing())
            minspacing = min(spacing)
            samp = []
            for s in spacing:
                samp.append(s/minspacing)
            if( debug ):
                print "sampling is ", samp
        elif( type(samp) == type(1) or type(samp) == type(1.0) ):
                samp = [samp]*numDim
        inputSpacing = img.GetSpacing()
        inputSize = img.GetLargestPossibleRegion().GetSize()
        if( debug ): print "selecting interoplator"
        if( interp.lower() == 'linear' ):
                interpolator = itk.LinearInterpolateImageFunction[imgIType,itk.D].New()
        elif( interp.lower() == 'bspline' ):
                interepolator = itk.BSplineInterpolateImageFunction[imgIType,itk.D].New()
        else:
                interpolator = itk.NearestNeighborInterpolateImageFunction[imgIType,itk.D].New()
        transform = itk.IdentityTransform[itk.D,numDim].New()
        
        outputSize = []
        outputSpacing = []
        for i in range(numDim ):   
                outputSize.append(int( inputSize.GetElement(i)*float(samp[i]) +0.5 ) )
                outputSpacing.append(inputSpacing.GetElement(i)/samp[i])
        
        
        if( debug ): print "setting interpolator"
        transform.SetIdentity()
        
        resampler.SetInterpolator(interpolator)
        resampler.SetTransform(transform)
        
        
        if( debug ): print "setting output spacing and size"
        resampler.SetOutputSpacing(outputSpacing)
        resampler.SetOutputOrigin(img.GetOrigin())
        resampler.SetSize(outputSize)
        img3 = resampler.GetOutput()
        img3.Update()
        if( returnSampling ):
            return img3, samp
        else:
            return img3

    except Exception, error:
        print "failed in subsampleImage", error
        
